package com.gitee.wsl.coroutines.strategy.race

import com.gitee.wsl.base.Or
import com.gitee.wsl.coroutines.api.PipeAble
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.async
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.selects.select
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext

/**
 * Races the participants [fa], [fb] on the provided [CoroutineContext].
 * The winner of the race cancels the other participants.
 * Cancelling the operation cancels all participants.
 *
 * Coroutine context is inherited from a [CoroutineScope], additional context elements can be specified with [ctx] argument.
 * If the combined context does not have any dispatcher nor any other [ContinuationInterceptor], then [Dispatchers.Default] is used.
 * **WARNING** If the combined context has a single threaded [ContinuationInterceptor], this function will not run [fa] & [fb] in parallel.
 *
 * ```kotlin
 * import arrow.core.Either
 * import arrow.fx.coroutines.*
 * import kotlinx.coroutines.Dispatchers
 * import kotlinx.coroutines.awaitCancellation
 *
 * suspend fun main(): Unit {
 *   suspend fun loser(): Int =
 *     guaranteeCase({ awaitCancellation() }) { exitCase ->
 *       println("I can never win the race. Finished with $exitCase.")
 *     }
 *
 *   val winner = raceN(Dispatchers.IO, { loser() }, { 5 })
 *
 *   val res = when(winner) {
 *     is Either.Left -> "Never always loses race"
 *     is Either.Right -> "Race was won with ${winner.value}"
 *   }
 *   //sampleEnd
 *   println(res)
 * }
 * ```
 * <!--- KNIT example-race2-02.kt -->
 *
 * @param fa task to participate in the race
 * @param fb task to participate in the race
 * @return either [Either.Left] if [fa] won the race, or [Either.Right] if [fb] won the race.
 * @see raceN for a function that ensures it runs in parallel on the [Dispatchers.Default].
 */
@OptIn(ExperimentalContracts::class)
@Suppress("LEAKED_IN_PLACE_LAMBDA")
suspend inline fun <A, B> PipeAble.Companion.raceN(
    ctx: CoroutineContext = EmptyCoroutineContext,
    crossinline fa: suspend CoroutineScope.() -> A,
    crossinline fb: suspend CoroutineScope.() -> B
): Or<A, B> {
    contract {
        callsInPlace(fa, InvocationKind.AT_MOST_ONCE)
        callsInPlace(fb, InvocationKind.AT_MOST_ONCE)
    }
    return coroutineScope {
        val a = async(ctx) { fa() }
        val b = async(ctx) { fb() }
        select<Or<A, B>> {
            a.onAwait.invoke { Or.One(it) }
            b.onAwait.invoke { Or.Two(it) }
        }.also {
            when (it) {
                is Or.One -> b.cancelAndJoin()
                is Or.Two -> a.cancelAndJoin()
            }
        }
    }
}

